{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "j0a4mTk9o1Qg",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Copyright 2019 Google Inc.\n",
    "\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dCpvgG0vwXAZ"
   },
   "source": [
    "#Predicting Movie Review Sentiment with BERT on TF Hub"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xiYrZKaHwV81"
   },
   "source": [
    "If you’ve been following Natural Language Processing over the past year, you’ve probably heard of BERT: Bidirectional Encoder Representations from Transformers. It’s a neural network architecture designed by Google researchers that’s totally transformed what’s state-of-the-art for NLP tasks, like text classification, translation, summarization, and question answering.\n",
    "\n",
    "Now that BERT's been added to [TF Hub](https://www.tensorflow.org/hub) as a loadable module, it's easy(ish) to add into existing Tensorflow text pipelines. In an existing pipeline, BERT can replace text embedding layers like ELMO and GloVE. Alternatively, [finetuning](http://wiki.fast.ai/index.php/Fine_tuning) BERT can provide both an accuracy boost and faster training time in many cases.\n",
    "\n",
    "Here, we'll train a model to predict whether an IMDB movie review is positive or negative using BERT in Tensorflow with tf hub. Some code was adapted from [this colab notebook](https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/tools/colab/bert_finetuning_with_cloud_tpus.ipynb). Let's get started!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hsZvic2YxnTz",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'tensorflow_hub'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-4-aeac30832f18>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mpandas\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mpd\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mtensorflow\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 4\u001b[1;33m \u001b[1;32mimport\u001b[0m \u001b[0mtensorflow_hub\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mhub\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      5\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mdatetime\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mdatetime\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      6\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'tensorflow_hub'"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cp5wfXDx5SPH"
   },
   "source": [
    "In addition to the standard libraries we imported above, we'll need to install BERT's python package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "jviywGyWyKsA",
    "outputId": "166f3005-d219-404f-b201-2a0b75480360",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "!pip install bert-tensorflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hhbGEfwgdEtw",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import bert\n",
    "from bert import run_classifier\n",
    "from bert import optimization\n",
    "from bert import tokenization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KVB3eOcjxxm1"
   },
   "source": [
    "Below, we'll set an output directory location to store our model output and checkpoints. This can be a local directory, in which case you'd set OUTPUT_DIR to the name of the directory you'd like to create. If you're running this code in Google's hosted Colab, the directory won't persist after the Colab session ends.\n",
    "\n",
    "Alternatively, if you're a GCP user, you can store output in a GCP bucket. To do that, set a directory name in OUTPUT_DIR and the name of the GCP bucket in the BUCKET field.\n",
    "\n",
    "Set DO_DELETE to rewrite the OUTPUT_DIR if it exists. Otherwise, Tensorflow will load existing model checkpoints from that directory (if they exist)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "US_EAnICvP7f",
    "outputId": "7780a032-31d4-4794-e6aa-664a5d2ae7dd",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Set the output directory for saving model file\n",
    "# Optionally, set a GCP bucket location\n",
    "\n",
    "OUTPUT_DIR = 'OUTPUT_DIR_NAME'#@param {type:\"string\"}\n",
    "#@markdown Whether or not to clear/delete the directory and create a new one\n",
    "DO_DELETE = False #@param {type:\"boolean\"}\n",
    "#@markdown Set USE_BUCKET and BUCKET if you want to (optionally) store model output on GCP bucket.\n",
    "USE_BUCKET = True #@param {type:\"boolean\"}\n",
    "BUCKET = 'BUCKET_NAME' #@param {type:\"string\"}\n",
    "\n",
    "if USE_BUCKET:\n",
    "  OUTPUT_DIR = 'gs://{}/{}'.format(BUCKET, OUTPUT_DIR)\n",
    "  from google.colab import auth\n",
    "  auth.authenticate_user()\n",
    "\n",
    "if DO_DELETE:\n",
    "  try:\n",
    "    tf.gfile.DeleteRecursively(OUTPUT_DIR)\n",
    "  except:\n",
    "    # Doesn't matter if the directory didn't exist\n",
    "    pass\n",
    "tf.gfile.MakeDirs(OUTPUT_DIR)\n",
    "print('***** Model output directory: {} *****'.format(OUTPUT_DIR))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pmFYvkylMwXn"
   },
   "source": [
    "#Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MC_w8SRqN0fr"
   },
   "source": [
    "First, let's download the dataset, hosted by Stanford. The code below, which downloads, extracts, and imports the IMDB Large Movie Review Dataset, is borrowed from [this Tensorflow tutorial](https://www.tensorflow.org/hub/tutorials/text_classification_with_tf_hub)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fom_ff20gyy6",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "import os\n",
    "import re\n",
    "\n",
    "# Load all files from a directory in a DataFrame.\n",
    "def load_directory_data(directory):\n",
    "  data = {}\n",
    "  data[\"sentence\"] = []\n",
    "  data[\"sentiment\"] = []\n",
    "  for file_path in os.listdir(directory):\n",
    "    with tf.gfile.GFile(os.path.join(directory, file_path), \"r\") as f:\n",
    "      data[\"sentence\"].append(f.read())\n",
    "      data[\"sentiment\"].append(re.match(\"\\d+_(\\d+)\\.txt\", file_path).group(1))\n",
    "  return pd.DataFrame.from_dict(data)\n",
    "\n",
    "# Merge positive and negative examples, add a polarity column and shuffle.\n",
    "def load_dataset(directory):\n",
    "  pos_df = load_directory_data(os.path.join(directory, \"pos\"))\n",
    "  neg_df = load_directory_data(os.path.join(directory, \"neg\"))\n",
    "  pos_df[\"polarity\"] = 1\n",
    "  neg_df[\"polarity\"] = 0\n",
    "  return pd.concat([pos_df, neg_df]).sample(frac=1).reset_index(drop=True)\n",
    "\n",
    "# Download and process the dataset files.\n",
    "def download_and_load_datasets(force_download=False):\n",
    "  dataset = tf.keras.utils.get_file(\n",
    "      fname=\"aclImdb.tar.gz\", \n",
    "      origin=\"http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\", \n",
    "      extract=True)\n",
    "  \n",
    "  train_df = load_dataset(os.path.join(os.path.dirname(dataset), \n",
    "                                       \"aclImdb\", \"train\"))\n",
    "  test_df = load_dataset(os.path.join(os.path.dirname(dataset), \n",
    "                                      \"aclImdb\", \"test\"))\n",
    "  \n",
    "  return train_df, test_df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2abfwdn-g135",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "train, test = download_and_load_datasets()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XA8WHJgzhIZf"
   },
   "source": [
    "To keep training fast, we'll take a sample of 5000 train and test examples, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lw_F488eixTV",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "train = train.sample(5000)\n",
    "test = test.sample(5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "prRQM8pDi8xI",
    "outputId": "34445cb8-2be0-4379-fdbc-7794091f6049",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "train.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "sfRnHSz3iSXz"
   },
   "source": [
    "For us, our input data is the 'sentence' column and our label is the 'polarity' column (0, 1 for negative and positive, respecitvely)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IuMOGwFui4it",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "DATA_COLUMN = 'sentence'\n",
    "LABEL_COLUMN = 'polarity'\n",
    "# label_list is the list of labels, i.e. True, False or 0, 1 or 'dog', 'cat'\n",
    "label_list = [0, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "V399W0rqNJ-Z"
   },
   "source": [
    "#Data Preprocessing\n",
    "We'll need to transform our data into a format BERT understands. This involves two steps. First, we create  `InputExample`'s using the constructor provided in the BERT library.\n",
    "\n",
    "- `text_a` is the text we want to classify, which in this case, is the `Request` field in our Dataframe. \n",
    "- `text_b` is used if we're training a model to understand the relationship between sentences (i.e. is `text_b` a translation of `text_a`? Is `text_b` an answer to the question asked by `text_a`?). This doesn't apply to our task, so we can leave `text_b` blank.\n",
    "- `label` is the label for our example, i.e. True, False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "p9gEt5SmM6i6",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'train' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-6-5445b9bc690e>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;31m# Use the InputExample class from BERT's run_classifier code to create examples from the data\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m train_InputExamples = train.apply(lambda x: bert.run_classifier.InputExample(guid=None, # Globally unique ID for bookkeeping, unused in this example\n\u001b[0m\u001b[0;32m      3\u001b[0m                                                                    \u001b[0mtext_a\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mx\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mDATA_COLUMN\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m                                                                    \u001b[0mtext_b\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m                                                                    label = x[LABEL_COLUMN]), axis = 1)\n",
      "\u001b[1;31mNameError\u001b[0m: name 'train' is not defined"
     ]
    }
   ],
   "source": [
    "# Use the InputExample class from BERT's run_classifier code to create examples from the data\n",
    "train_InputExamples = train.apply(lambda x: bert.run_classifier.InputExample(guid=None, # Globally unique ID for bookkeeping, unused in this example\n",
    "                                                                   text_a = x[DATA_COLUMN], \n",
    "                                                                   text_b = None, \n",
    "                                                                   label = x[LABEL_COLUMN]), axis = 1)\n",
    "\n",
    "test_InputExamples = test.apply(lambda x: bert.run_classifier.InputExample(guid=None, \n",
    "                                                                   text_a = x[DATA_COLUMN], \n",
    "                                                                   text_b = None, \n",
    "                                                                   label = x[LABEL_COLUMN]), axis = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SCZWZtKxObjh"
   },
   "source": [
    "Next, we need to preprocess our data so that it matches the data BERT was trained on. For this, we'll need to do a couple of things (but don't worry--this is also included in the Python library):\n",
    "\n",
    "\n",
    "1. Lowercase our text (if we're using a BERT lowercase model)\n",
    "2. Tokenize it (i.e. \"sally says hi\" -> [\"sally\", \"says\", \"hi\"])\n",
    "3. Break words into WordPieces (i.e. \"calling\" -> [\"call\", \"##ing\"])\n",
    "4. Map our words to indexes using a vocab file that BERT provides\n",
    "5. Add special \"CLS\" and \"SEP\" tokens (see the [readme](https://github.com/google-research/bert))\n",
    "6. Append \"index\" and \"segment\" tokens to each input (see the [BERT paper](https://arxiv.org/pdf/1810.04805.pdf))\n",
    "\n",
    "Happily, we don't have to worry about most of these details.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qMWiDtpyQSoU"
   },
   "source": [
    "To start, we'll need to load a vocabulary file and lowercasing information directly from the BERT tf hub module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "IhJSe0QHNG7U",
    "outputId": "20b28cc7-3cb3-4ce6-bfff-a7847ce3bbaa",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# This is a path to an uncased (all lowercase) version of BERT\n",
    "BERT_MODEL_HUB = \"https://tfhub.dev/google/bert_uncased_L-12_H-768_A-12/1\"\n",
    "\n",
    "def create_tokenizer_from_hub_module():\n",
    "  \"\"\"Get the vocab file and casing info from the Hub module.\"\"\"\n",
    "  with tf.Graph().as_default():\n",
    "    bert_module = hub.Module(BERT_MODEL_HUB)\n",
    "    tokenization_info = bert_module(signature=\"tokenization_info\", as_dict=True)\n",
    "    with tf.Session() as sess:\n",
    "      vocab_file, do_lower_case = sess.run([tokenization_info[\"vocab_file\"],\n",
    "                                            tokenization_info[\"do_lower_case\"]])\n",
    "      \n",
    "  return bert.tokenization.FullTokenizer(\n",
    "      vocab_file=vocab_file, do_lower_case=do_lower_case)\n",
    "\n",
    "tokenizer = create_tokenizer_from_hub_module()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "z4oFkhpZBDKm"
   },
   "source": [
    "Great--we just learned that the BERT model we're using expects lowercase data (that's what stored in tokenization_info[\"do_lower_case\"]) and we also loaded BERT's vocab file. We also created a tokenizer, which breaks words into word pieces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 221
    },
    "colab_type": "code",
    "id": "dsBo6RCtQmwx",
    "outputId": "9af8c917-90ec-4fe9-897b-79dc89ca88e1",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "tokenizer.tokenize(\"This here's an example of using the BERT tokenizer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0OEzfFIt6GIc"
   },
   "source": [
    "Using our tokenizer, we'll call `run_classifier.convert_examples_to_features` on our InputExamples to convert them into features BERT understands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1261
    },
    "colab_type": "code",
    "id": "LL5W8gEGRTAf",
    "outputId": "65001dda-155b-48fc-b5fc-1e4cabc8dfbf",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# We'll set sequences to be at most 128 tokens long.\n",
    "MAX_SEQ_LENGTH = 128\n",
    "# Convert our train and test features to InputFeatures that BERT understands.\n",
    "train_features = bert.run_classifier.convert_examples_to_features(train_InputExamples, label_list, MAX_SEQ_LENGTH, tokenizer)\n",
    "test_features = bert.run_classifier.convert_examples_to_features(test_InputExamples, label_list, MAX_SEQ_LENGTH, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ccp5trMwRtmr"
   },
   "source": [
    "#Creating a model\n",
    "\n",
    "Now that we've prepared our data, let's focus on building a model. `create_model` does just this below. First, it loads the BERT tf hub module again (this time to extract the computation graph). Next, it creates a single new layer that will be trained to adapt BERT to our sentiment task (i.e. classifying whether a movie review is positive or negative). This strategy of using a mostly trained model is called [fine-tuning](http://wiki.fast.ai/index.php/Fine_tuning)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6o2a5ZIvRcJq",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def create_model(is_predicting, input_ids, input_mask, segment_ids, labels,\n",
    "                 num_labels):\n",
    "  \"\"\"Creates a classification model.\"\"\"\n",
    "\n",
    "  bert_module = hub.Module(\n",
    "      BERT_MODEL_HUB,\n",
    "      trainable=True)\n",
    "  bert_inputs = dict(\n",
    "      input_ids=input_ids,\n",
    "      input_mask=input_mask,\n",
    "      segment_ids=segment_ids)\n",
    "  bert_outputs = bert_module(\n",
    "      inputs=bert_inputs,\n",
    "      signature=\"tokens\",\n",
    "      as_dict=True)\n",
    "\n",
    "  # Use \"pooled_output\" for classification tasks on an entire sentence.\n",
    "  # Use \"sequence_outputs\" for token-level output.\n",
    "  output_layer = bert_outputs[\"pooled_output\"]\n",
    "\n",
    "  hidden_size = output_layer.shape[-1].value\n",
    "\n",
    "  # Create our own layer to tune for politeness data.\n",
    "  output_weights = tf.get_variable(\n",
    "      \"output_weights\", [num_labels, hidden_size],\n",
    "      initializer=tf.truncated_normal_initializer(stddev=0.02))\n",
    "\n",
    "  output_bias = tf.get_variable(\n",
    "      \"output_bias\", [num_labels], initializer=tf.zeros_initializer())\n",
    "\n",
    "  with tf.variable_scope(\"loss\"):\n",
    "\n",
    "    # Dropout helps prevent overfitting\n",
    "    output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)\n",
    "\n",
    "    logits = tf.matmul(output_layer, output_weights, transpose_b=True)\n",
    "    logits = tf.nn.bias_add(logits, output_bias)\n",
    "    log_probs = tf.nn.log_softmax(logits, axis=-1)\n",
    "\n",
    "    # Convert labels into one-hot encoding\n",
    "    one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)\n",
    "\n",
    "    predicted_labels = tf.squeeze(tf.argmax(log_probs, axis=-1, output_type=tf.int32))\n",
    "    # If we're predicting, we want predicted labels and the probabiltiies.\n",
    "    if is_predicting:\n",
    "      return (predicted_labels, log_probs)\n",
    "\n",
    "    # If we're train/eval, compute loss between predicted and actual label\n",
    "    per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)\n",
    "    loss = tf.reduce_mean(per_example_loss)\n",
    "    return (loss, predicted_labels, log_probs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qpE0ZIDOCQzE"
   },
   "source": [
    "Next we'll wrap our model function in a `model_fn_builder` function that adapts our model to work for training, evaluation, and prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FnH-AnOQ9KKW",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# model_fn_builder actually creates our model function\n",
    "# using the passed parameters for num_labels, learning_rate, etc.\n",
    "def model_fn_builder(num_labels, learning_rate, num_train_steps,\n",
    "                     num_warmup_steps):\n",
    "  \"\"\"Returns `model_fn` closure for TPUEstimator.\"\"\"\n",
    "  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument\n",
    "    \"\"\"The `model_fn` for TPUEstimator.\"\"\"\n",
    "\n",
    "    input_ids = features[\"input_ids\"]\n",
    "    input_mask = features[\"input_mask\"]\n",
    "    segment_ids = features[\"segment_ids\"]\n",
    "    label_ids = features[\"label_ids\"]\n",
    "\n",
    "    is_predicting = (mode == tf.estimator.ModeKeys.PREDICT)\n",
    "    \n",
    "    # TRAIN and EVAL\n",
    "    if not is_predicting:\n",
    "\n",
    "      (loss, predicted_labels, log_probs) = create_model(\n",
    "        is_predicting, input_ids, input_mask, segment_ids, label_ids, num_labels)\n",
    "\n",
    "      train_op = bert.optimization.create_optimizer(\n",
    "          loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu=False)\n",
    "\n",
    "      # Calculate evaluation metrics. \n",
    "      def metric_fn(label_ids, predicted_labels):\n",
    "        accuracy = tf.metrics.accuracy(label_ids, predicted_labels)\n",
    "        f1_score = tf.contrib.metrics.f1_score(\n",
    "            label_ids,\n",
    "            predicted_labels)\n",
    "        auc = tf.metrics.auc(\n",
    "            label_ids,\n",
    "            predicted_labels)\n",
    "        recall = tf.metrics.recall(\n",
    "            label_ids,\n",
    "            predicted_labels)\n",
    "        precision = tf.metrics.precision(\n",
    "            label_ids,\n",
    "            predicted_labels) \n",
    "        true_pos = tf.metrics.true_positives(\n",
    "            label_ids,\n",
    "            predicted_labels)\n",
    "        true_neg = tf.metrics.true_negatives(\n",
    "            label_ids,\n",
    "            predicted_labels)   \n",
    "        false_pos = tf.metrics.false_positives(\n",
    "            label_ids,\n",
    "            predicted_labels)  \n",
    "        false_neg = tf.metrics.false_negatives(\n",
    "            label_ids,\n",
    "            predicted_labels)\n",
    "        return {\n",
    "            \"eval_accuracy\": accuracy,\n",
    "            \"f1_score\": f1_score,\n",
    "            \"auc\": auc,\n",
    "            \"precision\": precision,\n",
    "            \"recall\": recall,\n",
    "            \"true_positives\": true_pos,\n",
    "            \"true_negatives\": true_neg,\n",
    "            \"false_positives\": false_pos,\n",
    "            \"false_negatives\": false_neg\n",
    "        }\n",
    "\n",
    "      eval_metrics = metric_fn(label_ids, predicted_labels)\n",
    "\n",
    "      if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        return tf.estimator.EstimatorSpec(mode=mode,\n",
    "          loss=loss,\n",
    "          train_op=train_op)\n",
    "      else:\n",
    "          return tf.estimator.EstimatorSpec(mode=mode,\n",
    "            loss=loss,\n",
    "            eval_metric_ops=eval_metrics)\n",
    "    else:\n",
    "      (predicted_labels, log_probs) = create_model(\n",
    "        is_predicting, input_ids, input_mask, segment_ids, label_ids, num_labels)\n",
    "\n",
    "      predictions = {\n",
    "          'probabilities': log_probs,\n",
    "          'labels': predicted_labels\n",
    "      }\n",
    "      return tf.estimator.EstimatorSpec(mode, predictions=predictions)\n",
    "\n",
    "  # Return the actual model function in the closure\n",
    "  return model_fn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OjwJ4bTeWXD8",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Compute train and warmup steps from batch size\n",
    "# These hyperparameters are copied from this colab notebook (https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/tools/colab/bert_finetuning_with_cloud_tpus.ipynb)\n",
    "BATCH_SIZE = 32\n",
    "LEARNING_RATE = 2e-5\n",
    "NUM_TRAIN_EPOCHS = 3.0\n",
    "# Warmup is a period of time where hte learning rate \n",
    "# is small and gradually increases--usually helps training.\n",
    "WARMUP_PROPORTION = 0.1\n",
    "# Model configs\n",
    "SAVE_CHECKPOINTS_STEPS = 500\n",
    "SAVE_SUMMARY_STEPS = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "emHf9GhfWBZ_",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Compute # train and warmup steps from batch size\n",
    "num_train_steps = int(len(train_features) / BATCH_SIZE * NUM_TRAIN_EPOCHS)\n",
    "num_warmup_steps = int(num_train_steps * WARMUP_PROPORTION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oEJldMr3WYZa",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Specify outpit directory and number of checkpoint steps to save\n",
    "run_config = tf.estimator.RunConfig(\n",
    "    model_dir=OUTPUT_DIR,\n",
    "    save_summary_steps=SAVE_SUMMARY_STEPS,\n",
    "    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 156
    },
    "colab_type": "code",
    "id": "q_WebpS1X97v",
    "outputId": "1648932a-7391-49d3-8af7-52d514e226e8",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "model_fn = model_fn_builder(\n",
    "  num_labels=len(label_list),\n",
    "  learning_rate=LEARNING_RATE,\n",
    "  num_train_steps=num_train_steps,\n",
    "  num_warmup_steps=num_warmup_steps)\n",
    "\n",
    "estimator = tf.estimator.Estimator(\n",
    "  model_fn=model_fn,\n",
    "  config=run_config,\n",
    "  params={\"batch_size\": BATCH_SIZE})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NOO3RfG1DYLo"
   },
   "source": [
    "Next we create an input builder function that takes our training feature set (`train_features`) and produces a generator. This is a pretty standard design pattern for working with Tensorflow [Estimators](https://www.tensorflow.org/guide/estimators)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1Pv2bAlOX_-K",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Create an input function for training. drop_remainder = True for using TPUs.\n",
    "train_input_fn = bert.run_classifier.input_fn_builder(\n",
    "    features=train_features,\n",
    "    seq_length=MAX_SEQ_LENGTH,\n",
    "    is_training=True,\n",
    "    drop_remainder=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t6Nukby2EB6-"
   },
   "source": [
    "Now we train our model! For me, using a Colab notebook running on Google's GPUs, my training time was about 14 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "nucD4gluYJmK",
    "outputId": "5d728e72-4631-42bf-c48d-3f51d4b968ce",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "print(f'Beginning Training!')\n",
    "current_time = datetime.now()\n",
    "estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)\n",
    "print(\"Training took time \", datetime.now() - current_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CmbLTVniARy3"
   },
   "source": [
    "Now let's use our test data to see how well our model did:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JIhejfpyJ8Bx",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "test_input_fn = run_classifier.input_fn_builder(\n",
    "    features=test_features,\n",
    "    seq_length=MAX_SEQ_LENGTH,\n",
    "    is_training=False,\n",
    "    drop_remainder=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 445
    },
    "colab_type": "code",
    "id": "PPVEXhNjYXC-",
    "outputId": "dd5482cd-c558-465f-c854-ec11a0175316",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "estimator.evaluate(input_fn=test_input_fn, steps=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ueKsULteiz1B"
   },
   "source": [
    "Now let's write code to make predictions on new sentences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OsrbTD2EJTVl",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def getPrediction(in_sentences):\n",
    "  labels = [\"Negative\", \"Positive\"]\n",
    "  input_examples = [run_classifier.InputExample(guid=\"\", text_a = x, text_b = None, label = 0) for x in in_sentences] # here, \"\" is just a dummy label\n",
    "  input_features = run_classifier.convert_examples_to_features(input_examples, label_list, MAX_SEQ_LENGTH, tokenizer)\n",
    "  predict_input_fn = run_classifier.input_fn_builder(features=input_features, seq_length=MAX_SEQ_LENGTH, is_training=False, drop_remainder=False)\n",
    "  predictions = estimator.predict(predict_input_fn)\n",
    "  return [(sentence, prediction['probabilities'], labels[prediction['labels']]) for sentence, prediction in zip(in_sentences, predictions)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-thbodgih_VJ",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "pred_sentences = [\n",
    "  \"That movie was absolutely awful\",\n",
    "  \"The acting was a bit lacking\",\n",
    "  \"The film was creative and surprising\",\n",
    "  \"Absolutely fantastic!\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 649
    },
    "colab_type": "code",
    "id": "QrZmvZySKQTm",
    "outputId": "3891fafb-a460-4eb8-fa6c-335a5bbc10e5",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "predictions = getPrediction(pred_sentences)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MXkRiEBUqN3n"
   },
   "source": [
    "Voila! We have a sentiment classifier!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 221
    },
    "colab_type": "code",
    "id": "ERkTE8-7oQLZ",
    "outputId": "26c33224-dc2c-4b3d-f7b4-ac3ef0a58b27",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "predictions"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Predicting Movie Reviews with BERT on TF Hub.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
